--- external/eg3d/training/triplane.py	2023-03-31 20:15:14.373619528 +0000
+++ external_reference/eg3d/training/triplane.py	2023-03-31 20:14:21.579907466 +0000
@@ -10,11 +10,13 @@
 
 import torch
 from torch_utils import persistence
-from training.networks_stylegan2 import Generator as StyleGAN2Backbone
-from training.volumetric_rendering.renderer import ImportanceRenderer
-from training.volumetric_rendering.ray_sampler import RaySampler
+from external.stylegan.training.networks_stylegan2_terrain import Generator as StyleGAN2Backbone
+from external.eg3d.training.volumetric_rendering.renderer import ImportanceRenderer
+from external.eg3d.training.volumetric_rendering.ray_sampler import RaySampler
 import dnnlib
 
+from utils import noise_util
+
 @persistence.persistent_class
 class TriPlaneGenerator(torch.nn.Module):
     def __init__(self,
@@ -23,6 +25,7 @@
         w_dim,                      # Intermediate latent (W) dimensionality.
         img_resolution,             # Output resolution.
         img_channels,               # Number of output color channels.
+        plane_resolution    = 256,  # Resolution of feature planes 
         sr_num_fp16_res     = 0,
         mapping_kwargs      = {},   # Arguments for MappingNetwork.
         rendering_kwargs    = {},
@@ -35,22 +38,33 @@
         self.w_dim=w_dim
         self.img_resolution=img_resolution
         self.img_channels=img_channels
+        self.plane_resolution=plane_resolution
         self.renderer = ImportanceRenderer()
         self.ray_sampler = RaySampler()
-        self.backbone = StyleGAN2Backbone(z_dim, c_dim, w_dim, img_resolution=256, img_channels=32*3, mapping_kwargs=mapping_kwargs, **synthesis_kwargs)
+        self.backbone_xy = StyleGAN2Backbone(z_dim, c_dim, w_dim, img_resolution=plane_resolution, img_channels=32, mapping_kwargs=mapping_kwargs, **synthesis_kwargs)
+        self.backbone_xz = StyleGAN2Backbone(z_dim, c_dim, w_dim, img_resolution=plane_resolution, img_channels=32, mapping_kwargs=mapping_kwargs, **synthesis_kwargs)
+        self.backbone_yz = StyleGAN2Backbone(z_dim, c_dim, w_dim, img_resolution=plane_resolution, img_channels=32, mapping_kwargs=mapping_kwargs, **synthesis_kwargs)
+        # mapping network is only from backbone_xy
+        self.backbone_xz.mapping = None
+        self.backbone_yz.mapping = None
         self.superresolution = dnnlib.util.construct_class_by_name(class_name=rendering_kwargs['superresolution_module'], channels=32, img_resolution=img_resolution, sr_num_fp16_res=sr_num_fp16_res, sr_antialias=rendering_kwargs['sr_antialias'], **sr_kwargs)
         self.decoder = OSGDecoder(32, {'decoder_lr_mul': rendering_kwargs.get('decoder_lr_mul', 1), 'decoder_output_dim': 32})
         self.neural_rendering_resolution = 64
         self.rendering_kwargs = rendering_kwargs
-    
+
         self._last_planes = None
-    
+        self.noise_generator = noise_util.NoiseGenerator(
+            plane_resolution, plane_resolution, plane_resolution,
+            rendering_kwargs['box_warp'], rendering_kwargs['box_warp'], rendering_kwargs['box_warp'])
+
     def mapping(self, z, c, truncation_psi=1, truncation_cutoff=None, update_emas=False):
         if self.rendering_kwargs['c_gen_conditioning_zero']:
                 c = torch.zeros_like(c)
-        return self.backbone.mapping(z, c * self.rendering_kwargs.get('c_scale', 0), truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff, update_emas=update_emas)
+        return self.backbone_xy.mapping(z, c * self.rendering_kwargs.get('c_scale', 0), truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff, update_emas=update_emas)
 
-    def synthesis(self, ws, c, neural_rendering_resolution=None, update_emas=False, cache_backbone=False, use_cached_backbone=False, **synthesis_kwargs):
+    def synthesis(self, ws, c, neural_rendering_resolution=None,
+                  update_emas=False, cache_backbone=False, use_cached_backbone=False,
+                  noise_input=None, **synthesis_kwargs):
         cam2world_matrix = c[:, :16].view(-1, 4, 4)
         intrinsics = c[:, 16:25].view(-1, 3, 3)
 
@@ -67,38 +81,67 @@
         if use_cached_backbone and self._last_planes is not None:
             planes = self._last_planes
         else:
-            planes = self.backbone.synthesis(ws, update_emas=update_emas, **synthesis_kwargs)
+            plane_xy = self.backbone_xy.synthesis(ws, update_emas=update_emas, **synthesis_kwargs)
+            plane_xz = self.backbone_xz.synthesis(ws, update_emas=update_emas, **synthesis_kwargs)
+            plane_yz = self.backbone_yz.synthesis(ws, update_emas=update_emas, **synthesis_kwargs)
+            planes = torch.stack([plane_xy, plane_xz, plane_yz], dim=1) # N x 3 x 32 x H x W
         if cache_backbone:
             self._last_planes = planes
 
         # Reshape output into three 32-channel planes
-        planes = planes.view(len(planes), 3, 32, planes.shape[-2], planes.shape[-1])
+        # planes = planes.view(len(planes), 3, 32, planes.shape[-2], planes.shape[-1])
+
+        if noise_input is None:
+            noise_input = self.noise_generator.get_noise(ws.shape[0], ws.device)
 
         # Perform volume rendering
-        feature_samples, depth_samples, weights_samples = self.renderer(planes, self.decoder, ray_origins, ray_directions, self.rendering_kwargs) # channels last
+        feature_samples, depth_samples, disp_samples, weights_samples, noise_samples = self.renderer(planes, self.decoder, ray_origins, ray_directions, self.rendering_kwargs, noise_input=noise_input) # channels last
 
         # Reshape into 'raw' neural-rendered image
         H = W = self.neural_rendering_resolution
         feature_image = feature_samples.permute(0, 2, 1).reshape(N, feature_samples.shape[-1], H, W).contiguous()
         depth_image = depth_samples.permute(0, 2, 1).reshape(N, 1, H, W)
+        disp_image = disp_samples.permute(0, 2, 1).reshape(N, 1, H, W)
+        weights_image = weights_samples.permute(0, 2, 1).reshape(N, 1, H, W)
+        noise_image = noise_samples.permute(0, 2, 1).reshape(N, 1, H, W)
 
         # Run superresolution to get final image
         rgb_image = feature_image[:, :3]
-        sr_image = self.superresolution(rgb_image, feature_image, ws, noise_mode=self.rendering_kwargs['superresolution_noise_mode'], **{k:synthesis_kwargs[k] for k in synthesis_kwargs.keys() if k != 'noise_mode'})
-
-        return {'image': sr_image, 'image_raw': rgb_image, 'image_depth': depth_image}
+        rgb_and_disp_image = torch.cat([rgb_image, disp_image], dim=1)
+        # self.superresolution will modify rgb_and_disp_image in place
+        # if size=128 --> make a copy of it before input to SR network
+        sr_image_and_disp = self.superresolution(rgb_and_disp_image.clone(), feature_image, ws, noise_mode=self.rendering_kwargs['superresolution_noise_mode'], noise_input=noise_image, **{k:synthesis_kwargs[k] for k in synthesis_kwargs.keys() if k != 'noise_mode'})
+        sr_image = sr_image_and_disp[:, :3]
+        sr_disp = sr_image_and_disp[:, 3:]
+
+        return {'image': sr_image_and_disp,
+                'disp': sr_disp,
+                'image_raw': rgb_and_disp_image, # rgb_image,
+                'depth_raw': depth_image,
+                'disp_raw': disp_image,
+                'weights_raw': weights_image,
+                'noise_raw': noise_image
+               }
     
     def sample(self, coordinates, directions, z, c, truncation_psi=1, truncation_cutoff=None, update_emas=False, **synthesis_kwargs):
         # Compute RGB features, density for arbitrary 3D coordinates. Mostly used for extracting shapes. 
         ws = self.mapping(z, c, truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff, update_emas=update_emas)
-        planes = self.backbone.synthesis(ws, update_emas=update_emas, **synthesis_kwargs)
-        planes = planes.view(len(planes), 3, 32, planes.shape[-2], planes.shape[-1])
+        # planes = self.backbone.synthesis(ws, update_emas=update_emas, **synthesis_kwargs)
+        # planes = planes.view(len(planes), 3, 32, planes.shape[-2], planes.shape[-1])
+        plane_xy = self.backbone_xy.synthesis(ws, update_emas=update_emas, **synthesis_kwargs)
+        plane_xz = self.backbone_xz.synthesis(ws, update_emas=update_emas, **synthesis_kwargs)
+        plane_yz = self.backbone_yz.synthesis(ws, update_emas=update_emas, **synthesis_kwargs)
+        planes = torch.stack([plane_xy, plane_xz, plane_yz], dim=1) # N x 3 x 32 x H x W
         return self.renderer.run_model(planes, self.decoder, coordinates, directions, self.rendering_kwargs)
 
     def sample_mixed(self, coordinates, directions, ws, truncation_psi=1, truncation_cutoff=None, update_emas=False, **synthesis_kwargs):
         # Same as sample, but expects latent vectors 'ws' instead of Gaussian noise 'z'
-        planes = self.backbone.synthesis(ws, update_emas = update_emas, **synthesis_kwargs)
-        planes = planes.view(len(planes), 3, 32, planes.shape[-2], planes.shape[-1])
+        # planes = self.backbone.synthesis(ws, update_emas=update_emas, **synthesis_kwargs)
+        # planes = planes.view(len(planes), 3, 32, planes.shape[-2], planes.shape[-1])
+        plane_xy = self.backbone_xy.synthesis(ws, update_emas=update_emas, **synthesis_kwargs)
+        plane_xz = self.backbone_xz.synthesis(ws, update_emas=update_emas, **synthesis_kwargs)
+        plane_yz = self.backbone_yz.synthesis(ws, update_emas=update_emas, **synthesis_kwargs)
+        planes = torch.stack([plane_xy, plane_xz, plane_yz], dim=1) # N x 3 x 32 x H x W
         return self.renderer.run_model(planes, self.decoder, coordinates, directions, self.rendering_kwargs)
 
     def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, neural_rendering_resolution=None, update_emas=False, cache_backbone=False, use_cached_backbone=False, **synthesis_kwargs):
@@ -107,7 +150,7 @@
         return self.synthesis(ws, c, update_emas=update_emas, neural_rendering_resolution=neural_rendering_resolution, cache_backbone=cache_backbone, use_cached_backbone=use_cached_backbone, **synthesis_kwargs)
 
 
-from training.networks_stylegan2 import FullyConnectedLayer
+from external.stylegan.training.networks_stylegan2_terrain import FullyConnectedLayer
 
 class OSGDecoder(torch.nn.Module):
     def __init__(self, n_features, options):
